import os
import sys
from openpyxl import Workbook
from datetime import datetime
import csv
from typing import Dict, List, Tuple

def validate_arguments() -> Tuple[str, str]:
    if len(sys.argv) != 3:
        print("Usage: python3 script.py <csv_path> <analysis_xlsx_dir>")
        sys.exit(1)

    csv_path, analysis_xlsx_dir = sys.argv[1], sys.argv[2]

    if not os.path.isdir(csv_path):
        print(f"Error: {csv_path} is not a valid directory.")
        sys.exit(1)

    os.makedirs(analysis_xlsx_dir, exist_ok=True)

    return csv_path, analysis_xlsx_dir

def create_output_workbook(program_list: List[str], cpu_count: int) -> Workbook:
    new_workbook = Workbook()
    new_workbook.remove(new_workbook.active)  # Remove the default sheet

    # Create sheets for CPU (including per-core sheets), MEM, and each program
    cpu_sheets = ['CPU_ALL_RAW'] + [f'CPU_{i}_RAW' for i in range(cpu_count)]
    sheet_names = cpu_sheets + ['MEM_RAW'] + [f"{program.upper()}_RAW" for program in program_list]
    for sheet_name in sheet_names:
        new_workbook.create_sheet(title=sheet_name)
        print(f"Created sheet: {sheet_name}")

    return new_workbook

def process_cpu_usage_file(filepath: str, new_workbook: Workbook):
    print(f"Processing CPU usage file: {filepath}")
    cpu_sheets = {sheet_name: new_workbook[sheet_name] for sheet_name in new_workbook.sheetnames if sheet_name.startswith('CPU_')}

    with open(filepath, 'r') as file:
        csv_reader = csv.reader(file)
        lines = list(csv_reader)

        # Write testing date to all CPU sheets
        testing_date = lines[0][0] if lines[0] else "No testing date found"
        for sheet in cpu_sheets.values():
            sheet['A1'] = testing_date

        # Write header to all CPU sheets
        header = lines[2] if len(lines) > 2 else []
        for sheet in cpu_sheets.values():
            for col, value in enumerate(header, start=1):
                sheet.cell(row=3, column=col, value=value)

        # Process and write data
        for line in lines[3:]:
            if len(line) < 2:
                continue
            cpu = line[1]
            sheet_name = f'CPU_{cpu}_RAW' if cpu != 'all' else 'CPU_ALL_RAW'
            if sheet_name in cpu_sheets:
                sheet = cpu_sheets[sheet_name]
                row = sheet.max_row + 1
                for col, value in enumerate(line, start=1):
                    cell = sheet.cell(row=row, column=col)
                    try:
                        cell.value = float(value) if value.replace('.', '', 1).isdigit() else value
                    except ValueError:
                        cell.value = value

def process_csv_files(csv_path: str, new_workbook: Workbook, worksheet_mapping: Dict[str, str]):
    for filename in os.listdir(csv_path):
        if not filename.endswith('.csv'):
            continue

        filepath = os.path.join(csv_path, filename)

        if filename == 'cpu_usage.csv':
            process_cpu_usage_file(filepath, new_workbook)
            continue

        worksheet_name = worksheet_mapping.get(filename)
        if not worksheet_name:
            print(f"Warning: No worksheet mapping found for {filename}, skipping.")
            continue

        print(f"Processing {filename} -> {worksheet_name}")
        worksheet = new_workbook[worksheet_name]

        with open(filepath, 'r') as file:
            csv_reader = csv.reader(file)
            lines = list(csv_reader)

            # Write testing date
            worksheet['A1'] = lines[0][0] if lines[0] else "No testing date found"

            # Write header
            header = lines[2] if len(lines) > 2 else []
            for col, value in enumerate(header, start=1):
                worksheet.cell(row=3, column=col, value=value)

            # Write data
            for row, line in enumerate(lines[3:], start=4):
                for col, value in enumerate(line, start=1):
                    cell = worksheet.cell(row=row, column=col)
                    try:
                        cell.value = float(value) if value.replace('.', '', 1).isdigit() else value
                    except ValueError:
                        cell.value = value

def main():
    csv_path, analysis_xlsx_dir = validate_arguments()

    output_xlsx = os.path.join(analysis_xlsx_dir, f'resource_raw_{datetime.now().strftime("%Y_%m_%d_%H_%M")}.xlsx')
    print(f"Analysis output file: {output_xlsx}")

    # Dynamically determine the list of programs
    program_list = [f.split('_resource_usage.csv')[0] for f in os.listdir(csv_path) if f.endswith('_resource_usage.csv')]
    print(f"Detected programs: {program_list}")

    # Determine CPU count (assuming 8 cores, adjust if needed)
    cpu_count = 8

    new_workbook = create_output_workbook(program_list, cpu_count)

    worksheet_mapping = {
        'mem_usage.csv': 'MEM_RAW',
        'cpu_usage.csv': 'CPU_ALL_RAW',
    }

    # Add dynamic mappings for program-specific files
    for program in program_list:
        csv_filename = f"{program}_resource_usage.csv"
        worksheet_name = f"{program.upper()}_RAW"
        worksheet_mapping[csv_filename] = worksheet_name

    process_csv_files(csv_path, new_workbook, worksheet_mapping)

    new_workbook.save(output_xlsx)
    print(f"Excel file has been created and saved as {output_xlsx}\n")

if __name__ == "__main__":
    main()
